import torch
import torch.nn as nn
import torch.nn.functional as F
from models.swin_decoder import SwinDecoder


class CrossTaskDecoder(nn.Module):
    def __init__(self, p):
        super(CrossTaskDecoder, self).__init__()
        self.tasks = p.TASKS.NAMES
        self.tasks_num = len(self.tasks)
        self.swin_decoder = SwinDecoder(task_num=self.tasks_num)
        self.task_head = nn.ModuleDict()
        for t in self.tasks:
            self.task_head[t] = nn.Conv2d(in_channels=96, out_channels=p.TASKS.NUM_OUTPUT[t], kernel_size=1, bias=False)

    def forward(self, task_features):
        all_task_embeding = self.feat_merging(task_features)
        out_embeding, out_hw_shape = self.swin_decoder(all_task_embeding)
        task_feature = self.feat_spliting(out_embeding, out_hw_shape)
        task_feature = {t: F.interpolate(task_feature[t], scale_factor=4, mode='bilinear', align_corners=True) for t in self.tasks}
        task_feature = {t: self.task_head[t](task_feature[t]) for t in self.tasks}

        return task_feature

    def feat_merging(self, task_features):
        B, C, H, W = task_features['semseg'].shape
        ###############cat-flat check#####################
        all_task_features = torch.cat([task_features[t] for t in self.tasks], dim=3)  # B,C,H,TW
        all_task_embeding = all_task_features.view(B, C, -1).transpose(1, 2).contiguous()  # B,THW,C

        return all_task_embeding

    def feat_spliting(self, out_embeding, out_hw_shape):
        task_feature = {}
        B, L, C = out_embeding.shape
        out_feature = out_embeding.view(B, out_hw_shape[0], out_hw_shape[1], C).permute(0, 3, 1, 2).contiguous() # B,C,H,TW
        out_feature = out_feature.chunk(self.tasks_num, dim=3)
        for t, feature in zip(self.tasks, out_feature):
            task_feature[t] = feature
        return task_feature  # T,B,C,H,W


if __name__ == '__main__':
    from easydict import EasyDict as edict
    p = edict()
    p.TASKS = edict()
    p.TASKS.NAMES = ['seg', 'depth']
    p.TASKS.NUM_OUTPUT = {'seg': 5, 'depth': 1}
    model = CrossTaskDecoder(p)
    x = torch.rand(2, 768, 20, 20)
    output = model({'seg': x, 'depth': x})
    print()